import subprocess
from pesp4.IKEv1.IKEv1_client import *
from pesp4.IKEv2.IKEv2_client import *
from scapy.all import *
from operator import methodcaller
import colorama
from colorama import Fore
from aalpy.base import SUL
from pesp4 import enums
from pesp4.scripts.ssh import reset_libreswan, reset_strongswan

class IKESUL(SUL):
    def __init__(self, my_IP, target_IP, iface=None, version='v2', auth_mode=enums.AuthMethod.PSK, 
                object=None, dir = ''):
        super().__init__()
        self.my_IP = my_IP
        self.target_IP = target_IP
        self.iface = iface
        self.version = version
        self.auth_mode = auth_mode
        self.object = object
        self.IKE_client = None
        colorama.init(autoreset=True)
        
    def reset(self):
        if self.object == 'strongswan':
            reset_strongswan()
        elif self.object == 'libreswan':
            reset_libreswan()
        if self.version == 'v1':
            self.IKE_client = IKEv1_client(self.my_IP, self.target_IP, self.iface)
        else:
            self.IKE_client = IKEv2_Client(self.my_IP, self.target_IP, self.iface, auth_mode=self.auth_mode)
        return True

    def process_query(self, request):
        if self.version == 'v1':
            self.IKE_client.current_abstractinput = request
            result = str(methodcaller(default.func_dir[request])(self.IKE_client))
        elif self.version == 'v2':
            result = str(self.IKE_client.sendAndRecv(request))
        if 'Response.' in result:
            result = result.replace("Response.", "")
        if 'Notify.' in result:
            result = result.replace("Notify.", "")
        if result in ['No_child_SA','No_IKE_SA','Have_Rekeyed','Un_supported']:
            result = 'None'
        print(f'{request} | {result}')
        return result
        
    def pre(self):
        attempt_time = 0
        while attempt_time < 5:
            if self.reset():
                break
            attempt_time += 1
        if attempt_time == 5:
            print(Fore.RED + f'reset failed {attempt_time} times')
            sys.exit()
    
    def post(self):
        if self.object == 'winserver':
            delIKE = 'INFO_DelIKE' if self.version == 'v2' else 'delete_IKE'
            self.process_query(delIKE)
    
    def step(self, letter):
        letter = letter.replace('*', '')
        return self.process_query(letter)
    
    def query(self, word):
        self.pre()
        if len(word) == 0:
            out = [self.step(None)]
        else:
            out = [self.step(letter) for letter in word]
        # name = f'query-{self.num_queries}'
        # self.save_pcap(name)
        self.post()
        self.num_queries += 1
        self.num_steps += len(word)
        self.performed_steps_in_query = self.num_steps
        return out
    
